#pragma once

#include "vkinfer/runtime/vk/vkop.h"

namespace vkinfer
{
    class VkOpConcat : public VkOp
    {
    private:
        struct PushConstants
        {
            uint32_t dim1;
            uint32_t dim2;

            uint32_t input_b;
            uint32_t input_c;
            uint32_t input_h;
            uint32_t input_w;
            uint32_t outputput_b;
            uint32_t outputput_c;
            uint32_t outputput_h;
            uint32_t outputput_w;
            
            uint32_t axis = 0;
        } pc;

    public:
        VkOpConcat(VkDevice device, const std::shared_ptr<OnnxNodeInfo>& opinfo);

        virtual void update_descriptors(
            const std::vector<std::shared_ptr<VkBlob>>& inputs,
            std::vector<std::shared_ptr<VkBlob>>& outputs) override;
        virtual void forward(const std::vector<std::shared_ptr<Blob>>& inputs,
            std::vector<std::shared_ptr<Blob>>& outputs) override;
    };
}
